import asyncio
from src.database.chromadb_client import ChromaDBClient

\async def main():
    """验证 ChromaDBClient.store_document 与 search_documents 功能"""
    # 初始化客户端，使用项目默认数据目录
    client = ChromaDBClient(persist_directory="d:/MCP/mcp2/data/chroma_db")

    # 测试数据
    doc_id = "test_doc_1"
    content = "这是一个测试文档内容，用于验证写入功能。"
    metadata = {"user_id": "test_user"}

    # 1. 写入文档
    await client.store_document(doc_id, content, metadata)
    print(f"✅ 文档写入成功: {doc_id}")

    # 2. 搜索文档
    results = await client.search_documents(
        collection_name="documents",
        query="测试文档",
        n_results=5
    )
    print("搜索结果:", results)

    # 3. 校验写入与检索
    assert doc_id in results.get("ids", []), "❌ 写入的文档未在搜索结果中找到"
    print("✅ 搜索验证成功, 文档已检索到")


if __name__ == "__main__":
    asyncio.run(main())